{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# W&B Models\n",
    "## 1. Create an experiment"
   ],
   "id": "c3e4ee3e1fb1298c"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "import random\n",
    "\n",
    "from torchvision import models\n",
    "\n",
    "import wandb\n",
    "\n",
    "with wandb.init(\n",
    "        # Set the wandb entity where your project will be logged (generally your team name).\n",
    "        entity=\"lzz35-sau\",\n",
    "        # Set the wandb project where this run will be logged.\n",
    "        project=\"quickstart-project\",\n",
    "        notes=\"My first experiment\",\n",
    "        tags=[\"baseline\", \"paper1\"],\n",
    "        # Record the run's hyperparameters.\n",
    "        config={\"epochs\": 100, \"learning_rate\": 0.001, \"batch_size\": 64},\n",
    ") as run:\n",
    "    # Set up model and data.\n",
    "    model = models.resnet18(False)\n",
    "\n",
    "    # Run your training while logging metrics to visualize model performance.\n",
    "    offset = random.random() / 5\n",
    "    for epoch in range(1, 100):\n",
    "        acc = 1 - 2 ** -epoch - random.random() / epoch - offset\n",
    "        loss = 2 ** -epoch + random.random() / epoch + offset\n",
    "\n",
    "        # Log metrics to wandb.\n",
    "        run.log({\"acc\": acc, \"loss\": loss})\n",
    "\n",
    "    # Upload the trained model as an artifact.\n",
    "    # torch.save(model, \"res18.pth\")\n",
    "    # run.log_artifact(\"res18.pth\", name=\"trained-model\", type=\"model\")\n",
    "    run.finish()"
   ],
   "id": "3c8c7feb38a5457c"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## 2. Sweeps\n",
    "### 2.1 Tutorial: Define, initialize, and run a sweep"
   ],
   "id": "4b23921259b009d2"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "# Import the W&B Python Library and log into W&B\n",
    "import wandb\n",
    "\n",
    "wandb.login()\n",
    "\n",
    "\n",
    "# 1: Define objective/training function\n",
    "def objective(config):\n",
    "    score = config.x ** 3 + config.y\n",
    "    return score\n",
    "\n",
    "\n",
    "def main():\n",
    "    wandb.init(project=\"quickstart-project\")\n",
    "    score = objective(wandb.config)\n",
    "    wandb.log({\"score\": score})\n",
    "\n",
    "\n",
    "# 2: Define the search space\n",
    "sweep_configuration = {\n",
    "    \"method\": \"random\",\n",
    "    \"metric\": {\"goal\": \"minimize\", \"name\": \"score\"},\n",
    "    \"parameters\": {\n",
    "        \"x\": {\"max\": 0.1, \"min\": 0.01},\n",
    "        \"y\": {\"values\": [1, 3, 7]},\n",
    "    },\n",
    "}\n",
    "\n",
    "# 3: Start the sweep\n",
    "sweep_id = wandb.sweep(sweep=sweep_configuration, project=\"quickstart-project\")\n",
    "\n",
    "wandb.agent(sweep_id, function=main, count=10)\n"
   ],
   "id": "8331d52c76879097"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "### 2.2 Add W&B (wandb) to your code",
   "id": "e5c6263061e1cc53"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "import wandb\n",
    "import numpy as np\n",
    "import random\n",
    "\n",
    "# Define sweep config\n",
    "sweep_configuration = {\n",
    "    \"method\": \"random\",\n",
    "    \"name\": \"sweep\",\n",
    "    \"metric\": {\"goal\": \"maximize\", \"name\": \"val_acc\"},\n",
    "    \"parameters\": {\n",
    "        \"batch_size\": {\"values\": [16, 32, 64]},\n",
    "        \"epochs\": {\"values\": [5, 10, 15]},\n",
    "        \"lr\": {\"max\": 0.1, \"min\": 0.0001},\n",
    "    },\n",
    "}\n",
    "\n",
    "# Initialize sweep by passing in config.\n",
    "# (Optional) Provide a name of the project.\n",
    "sweep_id = wandb.sweep(sweep=sweep_configuration, project=\"quickstart-project\")\n",
    "\n",
    "\n",
    "# Define training function that takes in hyperparameter\n",
    "# values from `wandb.config` and uses them to train a\n",
    "# model and return metric\n",
    "def train_one_epoch(epoch, lr, bs):\n",
    "    acc = 0.25 + ((epoch / 30) + (random.random() / 10))\n",
    "    loss = 0.2 + (1 - ((epoch - 1) / 10 + random.random() / 5))\n",
    "    return acc, loss\n",
    "\n",
    "\n",
    "def evaluate_one_epoch(epoch):\n",
    "    acc = 0.1 + ((epoch / 20) + (random.random() / 10))\n",
    "    loss = 0.25 + (1 - ((epoch - 1) / 10 + random.random() / 6))\n",
    "    return acc, loss\n",
    "\n",
    "\n",
    "def main():\n",
    "    run = wandb.init()\n",
    "\n",
    "    # note that we define values from `wandb.config`\n",
    "    # instead of defining hard values\n",
    "    lr = wandb.config.lr\n",
    "    bs = wandb.config.batch_size\n",
    "    epochs = wandb.config.epochs\n",
    "\n",
    "    for epoch in np.arange(1, epochs):\n",
    "        train_acc, train_loss = train_one_epoch(epoch, lr, bs)\n",
    "        val_acc, val_loss = evaluate_one_epoch(epoch)\n",
    "\n",
    "        wandb.log(\n",
    "            {\n",
    "                \"epoch\": epoch,\n",
    "                \"train_acc\": train_acc,\n",
    "                \"train_loss\": train_loss,\n",
    "                \"val_acc\": val_acc,\n",
    "                \"val_loss\": val_loss,\n",
    "            }\n",
    "        )\n",
    "\n",
    "\n",
    "# Start sweep job.\n",
    "wandb.agent(sweep_id, function=main, count=4)\n"
   ],
   "id": "24d67edf53bbc1a7"
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}
